"""
随机采样并可视化训练数据和标签
用于检查数据生成是否正确
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
import os

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'Arial Unicode MS']
plt.rcParams['axes.unicode_minus'] = False

# 导入项目模块
import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from data_generator import SpectrumDataGenerator, SignalLabel
from signal_generator import SignalParams


def visualize_single_sample(spectrum, label, freq_axis, save_path=None):
    """可视化单个样本"""
    
    fig = plt.figure(figsize=(18, 12))
    gs = GridSpec(3, 3, figure=fig, hspace=0.3, wspace=0.3)
    
    # 1. 频谱图（全范围）
    ax1 = fig.add_subplot(gs[0, :])
    ax1.plot(freq_axis, spectrum, 'b-', linewidth=0.8, alpha=0.7, label='频谱')
    ax1.set_xlabel('频率 (Hz)', fontsize=12)
    ax1.set_ylabel('幅值', fontsize=12)
    ax1.set_title('频谱 (全频率范围)', fontsize=14, fontweight='bold')
    ax1.grid(True, alpha=0.3)
    ax1.legend()
    
    # 标注谐波位置
    if label.has_harmonic:
        for cluster in label.harmonic_clusters:
            for i in range(cluster['num_harmonics']):
                freq = cluster['base_freq'] * (i + 1)
                if freq < freq_axis[-1]:
                    ax1.axvline(freq, color='r', linestyle='--', alpha=0.5, linewidth=1,
                               label='谐波' if i == 0 else '')
    
    # 标注调制位置
    if label.has_modulation:
        for mod in label.modulation_info:
            carrier_freq = mod['carrier_freq']
            if carrier_freq < freq_axis[-1]:
                ax1.axvline(carrier_freq, color='g', linestyle=':', alpha=0.6, linewidth=1.5,
                           label='载波' if mod == label.modulation_info[0] else '')
                
                # 边带
                sideband_space = mod.get('sideband_spacing', mod['modulation_freq'])
                ax1.axvline(carrier_freq - sideband_space, color='orange', linestyle=':', 
                           alpha=0.4, linewidth=1, label='边带' if mod == label.modulation_info[0] else '')
                ax1.axvline(carrier_freq + sideband_space, color='orange', linestyle=':', 
                           alpha=0.4, linewidth=1)
    
    # 避免重复标签
    ax1.legend(loc='upper right', fontsize=9)
    
    # 2. 低频细节 (0-500 Hz)
    ax2 = fig.add_subplot(gs[1, 0])
    mask = freq_axis < 500
    ax2.plot(freq_axis[mask], spectrum[mask], 'b-', linewidth=1)
    ax2.set_xlabel('频率 (Hz)', fontsize=10)
    ax2.set_ylabel('幅值', fontsize=10)
    ax2.set_title('低频细节 (0-500 Hz)', fontsize=11, fontweight='bold')
    ax2.grid(True, alpha=0.3)
    
    # 标注谐波
    if label.has_harmonic:
        for cluster in label.harmonic_clusters:
            for i in range(cluster['num_harmonics']):
                freq = cluster['base_freq'] * (i + 1)
                if freq < 500:
                    ax2.axvline(freq, color='r', linestyle='--', alpha=0.6, linewidth=1.5)
    
    # 3. 中频细节 (500-1500 Hz)
    ax3 = fig.add_subplot(gs[1, 1])
    mask = (freq_axis >= 500) & (freq_axis < 1500)
    ax3.plot(freq_axis[mask], spectrum[mask], 'b-', linewidth=1)
    ax3.set_xlabel('频率 (Hz)', fontsize=10)
    ax3.set_ylabel('幅值', fontsize=10)
    ax3.set_title('中频细节 (500-1500 Hz)', fontsize=11, fontweight='bold')
    ax3.grid(True, alpha=0.3)
    
    # 标注调制
    if label.has_modulation:
        for mod in label.modulation_info:
            carrier_freq = mod['carrier_freq']
            if 500 <= carrier_freq < 1500:
                ax3.axvline(carrier_freq, color='g', linestyle=':', alpha=0.7, linewidth=1.5)
                sideband_space = mod.get('sideband_spacing', mod['modulation_freq'])
                ax3.axvline(carrier_freq - sideband_space, color='orange', linestyle=':', 
                           alpha=0.5, linewidth=1)
                ax3.axvline(carrier_freq + sideband_space, color='orange', linestyle=':', 
                           alpha=0.5, linewidth=1)
    
    # 4. 高频细节 (1500-2560 Hz)
    ax4 = fig.add_subplot(gs[1, 2])
    mask = (freq_axis >= 1500) & (freq_axis < 2560)
    ax4.plot(freq_axis[mask], spectrum[mask], 'b-', linewidth=1)
    ax4.set_xlabel('频率 (Hz)', fontsize=10)
    ax4.set_ylabel('幅值', fontsize=10)
    ax4.set_title('高频细节 (1500-2560 Hz)', fontsize=11, fontweight='bold')
    ax4.grid(True, alpha=0.3)
    
    # 5. 标签信息
    ax5 = fig.add_subplot(gs[2, :])
    ax5.axis('off')
    
    # 构建标签文本
    label_text = "=" * 80 + "\n"
    label_text += "标签信息\n"
    label_text += "=" * 80 + "\n\n"
    
    # 谐波信息
    label_text += "【谐波特征】\n"
    label_text += f"  检测: {'是' if label.has_harmonic else '否'}\n"
    if label.has_harmonic:
        label_text += f"  谐波簇数: {len(label.harmonic_clusters)}\n"
        for i, cluster in enumerate(label.harmonic_clusters):
            label_text += f"\n  簇 {i+1}:\n"
            label_text += f"    基频: {cluster['base_freq']:.2f} Hz\n"
            label_text += f"    谐波数: {cluster['num_harmonics']}\n"
            label_text += f"    谐波频率: {[cluster['base_freq']*(j+1) for j in range(cluster['num_harmonics'])]}\n"
            label_text += f"    幅值: {[f'{a:.3f}' for a in cluster['amplitudes']]}\n"
    label_text += "\n"
    
    # 调制信息
    label_text += "【调制特征】\n"
    label_text += f"  检测: {'是' if label.has_modulation else '否'}\n"
    if label.has_modulation:
        label_text += f"  调制数量: {len(label.modulation_info)}\n"
        for i, mod in enumerate(label.modulation_info):
            label_text += f"\n  调制 {i+1}:\n"
            label_text += f"    载波频率: {mod['carrier_freq']:.2f} Hz\n"
            label_text += f"    调制频率: {mod['modulation_freq']:.2f} Hz\n"
            label_text += f"    边带数量: {mod['num_sidebands']}\n"
            label_text += f"    边带间距: {mod.get('sideband_spacing', 0):.2f} Hz\n"
            label_text += f"    边带频率: [{mod['carrier_freq'] - mod.get('sideband_spacing', mod['modulation_freq']):.2f}, "
            label_text += f"{mod['carrier_freq'] + mod.get('sideband_spacing', mod['modulation_freq']):.2f}] Hz\n"
    label_text += "\n"
    
    # 轴承故障信息
    label_text += "【轴承故障】\n"
    label_text += f"  检测: {'是' if label.has_bearing_fault else '否'}\n"
    if label.has_bearing_fault and label.bearing_fault_info:
        info = label.bearing_fault_info
        label_text += f"  故障类型: {info.get('type', 'unknown')}\n"
        label_text += f"  故障频率: {info.get('fault_freq', 0):.2f} Hz\n"
        label_text += f"  共振频率: {info.get('resonance_freq', 0):.2f} Hz\n"
    label_text += "\n"
    
    # 统计信息
    label_text += "【频谱统计】\n"
    label_text += f"  峰值频率: {freq_axis[np.argmax(spectrum)]:.2f} Hz\n"
    label_text += f"  峰值幅值: {np.max(spectrum):.6f}\n"
    label_text += f"  平均值: {np.mean(spectrum):.6f}\n"
    label_text += f"  标准差: {np.std(spectrum):.6f}\n"
    
    ax5.text(0.05, 0.95, label_text, transform=ax5.transAxes,
             fontsize=10, verticalalignment='top', family='monospace',
             bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8))
    
    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        print(f"\n✓ 可视化已保存: {save_path}")
    else:
        plt.show()
    
    plt.close()
    return fig


def main():
    """主函数：生成并可视化随机样本"""
    
    print("=" * 80)
    print("训练数据可视化检查")
    print("=" * 80)
    
    # 创建数据生成器
    generator = SpectrumDataGenerator(fs=5120, n_samples=16384, seed=42)
    freq_axis = generator.get_freq_axis()
    
    print(f"\n数据生成器配置:")
    print(f"  采样频率: {generator.params.fs} Hz")
    print(f"  采样点数: {generator.params.n_samples}")
    print(f"  频谱长度: {len(freq_axis)}")
    print(f"  频率分辨率: {generator.params.fs / generator.params.n_samples:.4f} Hz")
    print(f"  Nyquist频率: {generator.max_freq} Hz")
    
    # 创建输出目录
    output_dir = "temp_scripts/visualization_output"
    os.makedirs(output_dir, exist_ok=True)
    
    # 生成5个不同类型的样本
    sample_types = [
        ("pure_harmonic", "纯谐波", generator.generate_pure_harmonic),
        ("modulation", "调制信号", generator.generate_modulation_signal),
        ("harmonic_modulation", "谐波+调制", generator.generate_harmonic_with_modulation),
        ("bearing_fault", "轴承故障", generator.generate_bearing_fault),
        ("complex", "复合信号", generator.generate_complex_signal),
    ]
    
    print(f"\n生成 {len(sample_types)} 个不同类型的样本...")
    print(f"输出目录: {output_dir}")
    
    for i, (name, desc, gen_func) in enumerate(sample_types):
        print(f"\n[{i+1}/{len(sample_types)}] 生成{desc}...")
        
        # 生成数据
        spectrum, label = gen_func()
        
        # 打印标签信息
        print(f"  ✓ 频谱长度: {len(spectrum)}")
        print(f"  ✓ 标签类型: {type(label).__name__}")
        print(f"    - 谐波: {label.has_harmonic}")
        print(f"    - 调制: {label.has_modulation}")
        print(f"    - 轴承故障: {label.has_bearing_fault}")
        
        # 可视化
        save_path = os.path.join(output_dir, f"{name}_sample.png")
        visualize_single_sample(spectrum, label, freq_axis, save_path)
    
    # 批量生成测试
    print(f"\n批量生成测试...")
    batch_size = 10
    spectra, labels = generator.generate_batch(batch_size)
    print(f"  ✓ 批量大小: {batch_size}")
    print(f"  ✓ 频谱形状: {spectra.shape}")
    print(f"  ✓ 标签数量: {len(labels)}")
    
    # 统计标签分布
    print(f"\n标签统计 (批量 {batch_size} 个样本):")
    harmonic_count = sum(1 for l in labels if l.has_harmonic)
    modulation_count = sum(1 for l in labels if l.has_modulation)
    bearing_count = sum(1 for l in labels if l.has_bearing_fault)
    
    print(f"  含谐波: {harmonic_count}/{batch_size} ({harmonic_count/batch_size*100:.1f}%)")
    print(f"  含调制: {modulation_count}/{batch_size} ({modulation_count/batch_size*100:.1f}%)")
    print(f"  含轴承故障: {bearing_count}/{batch_size} ({bearing_count/batch_size*100:.1f}%)")
    
    # 可视化随机3个批量样本
    print(f"\n可视化批量生成的随机3个样本...")
    for i in range(min(3, batch_size)):
        idx = i
        spectrum = spectra[idx]
        label = labels[idx]
        
        save_path = os.path.join(output_dir, f"batch_sample_{idx+1}.png")
        visualize_single_sample(spectrum, label, freq_axis, save_path)
    
    print("\n" + "=" * 80)
    print("✓ 所有可视化完成！")
    print("=" * 80)
    print(f"\n生成的文件:")
    for fname in sorted(os.listdir(output_dir)):
        if fname.endswith('.png'):
            print(f"  - {output_dir}/{fname}")
    
    print(f"\n建议检查:")
    print(f"  1. 频谱是否合理（没有异常噪声）")
    print(f"  2. 谐波频率是否标注正确")
    print(f"  3. 调制边带是否可见")
    print(f"  4. 标签信息是否完整准确")


if __name__ == "__main__":
    main()

